# depthbackbone网络类
from modules.depth_backbone import *
import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# device = torch.device("cpu")
class model(nn.Module):
    def __init__(self, n_proprio):
        super().__init__()
        self.depth_backbone = DepthOnlyFCBackbone58x87(32)
        self.depth_model = RecurrentDepthBackbone(self.depth_backbone, n_proprio)
        load_path = "/home/mi/MyRL/models/5500-vision_weight.pt"
        as_state_dict = torch.load(load_path)
        self.depth_model.to(device)
        self.depth_model.load_state_dict(as_state_dict['depth_encoder_state_dict'])